#include <Lucnx/interrupt.h>
#include <Lucnx/string.h>
#include <Lucnx/system.h>
#include <Lucnx/task.h>
#include <Lucnx/memory.h>
#include <Lucnx/debug.h>
#include <Lucnx/global.h>
#include <Lucnx/stdlib.h>
#include <Lucnx/printk.h>
#include <Lucnx/vsprintf.h>

#define MAX_PID 256
#define NR_TASKS 64

void task_switch(task_t *next);
void update_tss_esp(task_t *task);
void thread_init();

u8 pid_map[32];
static task_t *task_table[NR_TASKS];
extern task_t *idle_pcb;

void task_exec(void (*function)())
{
    set_interrupt_state(true);
    function();
    while (1);
}

task_t *running_task()
{
    // 取任务栈指针的整页部分就是任务的 PCB
    asm volatile(
        "movl %esp, %eax\n"
        "andl $0xFFFFF000, %eax\n");
}

// 分配一个任务 id
pid_t pid_alloc()
{
    for (size_t i = 0; i < 32; i++)
    {
        if (pid_map[i] == 0xFF)
        {
            continue;
        }
        for (size_t j = 0; j < 8; j++)
        {
            if (pid_map[i] & (1 << j))
            {
                continue;
            }
            pid_map[i] |= (1 << j); // 占位
            return i * 8 + j + 1;
        }
    }
    return EOF; // 找不到返回 EOF
}

// 释放任务 pid
void pid_free(task_t *task)
{
    pid_t pid = task->pid;
    int idx_by = pid / 8;
    int idx_bt = pid % 8;

    pid_map[idx_by] &= ~(1 << idx_bt);
}

// 在 task PCB 指针表里获取一个空闲位置
task_t *get_free_task()
{
    for (size_t i = 0; i < NR_TASKS; i++)
    {
        if (task_table[i] == NULL)
        {
            task_t *ptr = kpage_alloc(1); // 分配任务的 PCB 内核栈的内存地址
            task_table[i] = ptr;
            return task_table[i];
        }
    }
    return NULL;
}

extern task_t *idle_pcb;

// 搜索合适被调度并且状态为 state 的任务
task_t *search_task(task_state state)
{
    task_t *ptr;
    task_t *task;
    task_t *current = running_task();

    for (size_t i = 0; i < NR_TASKS; i++)
    {
        ptr = task_table[i];
        if (ptr == NULL || ptr == task || ptr == current || ptr->state != state)
        {
            continue;
        }
        if ((task == NULL) || (ptr->counter > task->counter) || (ptr->jiffies < task->jiffies))
        {
            task = ptr;
        }
    }

    // 找不到返回 idle 线程
    if (task == NULL)
    {
        return idle_pcb;
    }
    return task;
}

task_t *task_create(void *target, char *name, u8 prio)
{
    task_t *task = get_free_task();
    memset(task, 0, PAGE_SIZE); 

    task->kstack = (u32) task + PAGE_SIZE;
    task->kstack -= sizeof(interrupt_frame_t);
    task->kstack -= sizeof(task_frame_t); // 预留空间
    
    task_frame_t *task_frame = (task_frame_t*)task->kstack;
    task_frame->ebp = 0;
    task_frame->ebx = 0;
    task_frame->edi = 0;
    task_frame->esi = 0; 
    task_frame->eip = target;

    task->prio = prio;
    task->counter = prio;
    task->jiffies = 0;

    strcpy(task->name, name);
    task->state = TASK_READY;
    task->pid = pid_alloc();
    task->ppid = 0;
    task->magic = LUCNX_MAGIC;
    task->page_dir = NULL;
    task->mode = 1; // 默认为内核任务
    task->brk = KERNEL_MEM_SIZE;

    return task;
}

pid_t sys_getpid()
{
    return running_task()->pid;
}

pid_t sys_getppid()
{
    return running_task()->ppid;
}

void schedule()
{
    task_t *current = running_task();
    if (current->state == TASK_RUNNING)
    {
        current->state = TASK_READY;
        current->counter = current->prio;
    }

    task_t *next = search_task(TASK_READY);
    assert(next->magic == LUCNX_MAGIC); // 检测栈溢出

    next->state = TASK_RUNNING;
    task_activate(next);
    task_switch(next);
}

void task_init()
{
    memset(task_table, 0, sizeof(task_table));
    thread_init();
    INFOK("INFO", "Task init done");
}

void task_block(task_t *task, task_state state)
{
    set_interrupt_state(false);

    task->state = state;

    schedule();
    set_interrupt_state(true);
}

void task_unblock(task_t *task)
{
    bool state = set_interrupt_state(false);

    task->state = TASK_READY;
    task->counter = task->prio;

    set_interrupt_state(state);
}

static task_t *pid_to_task(pid_t ppid)
{
    return task_table[ppid - 1];
}

void task_activate(task_t *task)
{
    u32 page_dir_paddr = 0x1000;

    if (task->mode == 2)
    {
        page_dir_paddr = (u32)addr_v2p(task->page_dir);
        update_tss_esp(task);
    }

    asm volatile(
        "movl %%eax, %%cr3\n"::"a"(page_dir_paddr));
}

u32 *create_pde()
{
    u32 *pde = kpage_alloc(1);
    assert(pde != NULL);
    memcpy(pde, (u32*)(0xFFFFF000), 8);
    pde[1023] = ((u32)addr_v2p(pde)) | 7;

    return pde;
}

extern void interrupt_exit();
void task_to_user_mode(void *target)
{
    task_t *current = running_task();
    assert(current->mode == 1);
    current->mode = 2;

    current->page_dir = create_pde();
    current->brk = USER_MEMORY;
    task_activate(current);

    current->vmap.bits = kpage_alloc(1);
    current->vmap.btmp_len = PAGE_SIZE;
    bitmap_init(&current->vmap);

    build_task_stack3(); // 构建三级栈
    arena_init(current->desc_dir);
    interrupt_frame_t *iframe = (interrupt_frame_t*) ((u32)current + PAGE_SIZE - sizeof(interrupt_frame_t));

    iframe->eax = 0;
    iframe->ebx = 0;
    iframe->ecx = 0;
    iframe->edx = 0;
    iframe->esi = 0;
    iframe->edi = 0;
    iframe->ebp = 0;
    iframe->esp_dummy = 0;

    iframe->esp = USER_STACK_TOP;

    iframe->ds = SELECTOR_USER_DATA;
    iframe->es = SELECTOR_USER_DATA;
    iframe->fs = SELECTOR_USER_DATA;
    iframe->gs = SELECTOR_USER_DATA;
    iframe->ss = SELECTOR_USER_DATA;
    iframe->cs = SELECTOR_USER_CODE; // 用户态的段寄存器

    iframe->eflags = (0b10) | (1 << 9) | (0 << 12);

    iframe->eip = target;

    // 切换模式
    asm volatile(
        "movl %0, %%esp\n"
        "jmp _interrupt_exit\n"::"g"(iframe) : "memory");
}

// 任务主动让出 CPU
void task_yield()
{
    bool old_state = set_interrupt_state(false);
    schedule();
    set_interrupt_state(old_state);
}

extern u32 all_ticks;

void task_sleep(u32 ms)
{
    u32 ticks = ms / 10;
    u32 start_ticks = all_ticks;

    while (all_ticks - start_ticks < ticks)
    {
        task_yield();
    }
}

extern void copy_task_mem(task_t *task);
extern void free_task_mem();
extern void free_vmap();

pid_t sys_fork()
{
    task_t *task = get_free_task();
    task_t *parent = running_task();

    // 拷贝 PCB 和内核栈
    memcpy(task, parent, PAGE_SIZE);

    // 修改子进程信息
    task->counter = task->prio;
    task->state = TASK_READY;
    task->pid = pid_alloc();
    task->ppid = parent->pid;
    task->jiffies = 0;
    u32 kstack_in_page = parent->kstack & 0xFFF;
    task->kstack = (u32)task + kstack_in_page;
    task->task_node.next = NULL;
    task->task_node.prev = NULL;

    // 创建页目录
    task->page_dir = create_pde();

    // 创建虚拟内存位图
    task->vmap.bits = kpage_alloc(1);
    task->vmap.btmp_len = PAGE_SIZE;
    bitmap_init(&task->vmap);

    // 拷贝用户态内存和 3 级栈
    copy_task_mem(task);

    interrupt_frame_t *iframe = (interrupt_frame_t *) ((u32)task + PAGE_SIZE - sizeof(interrupt_frame_t));
    iframe->eax = 0;
    task_frame_t *frame = (task_frame_t *) ((u32)iframe - sizeof(task_frame_t));
    frame->eip = interrupt_exit; // 返回进入用户态 修改返回值
    frame->esi = 0;
    frame->edi = 0;
    frame->ebp = 0;
    task->kstack = (u32)frame;

    return task->pid;
}

// 任务结束
void sys_exit(int state)
{
    task_t *task = running_task();

    if (task->ppid == 0)
    {
        printk("exit(): error!! task '%s' hasn't parent\n");
        return;
    }
    task->exit_state = state;

    free_task_mem();
    free_pde();
    free_vmap();

    task_t *parent = pid_to_task(task->ppid);
    if (parent == NULL)
    {
        printk("exit() warning: can't find %s's parent %d\n", task->name, task->ppid);
        return;
    }

    // 将所有子进程过继给 test 进程
    for (size_t i = 0; i < NR_TASKS; i++)
    {
        task_t *child = task_table[i];
        if (!child || child->ppid != task->pid)
            continue;
        child->ppid = task->ppid;
    }

    if (parent->state == TASK_WAITING)
    {
        task_unblock(parent);
    }

    task_block(task, TASK_BLOCKED);
}

void ps()
{
    task_t *ptr;

    pad_print("NAME", 16);
    pad_print("PID", 16);
    pad_print("PPID", 16);
    pad_print("STATE", 16);
    pad_print("PTR", 16);
    putchk('\n');

    for (size_t i = 0; i < NR_TASKS; i++)
    {
        ptr = task_table[i];
        if (ptr == NULL)
        {
            continue;
        }
        
        pad_print(ptr->name, 16);

        char buf[8];
        sprintf(buf, "%d", ptr->pid);
        pad_print(buf, 16);

        memset(buf, 0, sizeof(buf));
        sprintf(buf, "%d", ptr->ppid);
        pad_print(buf, 16);

        if (ptr->state == TASK_RUNNING)
        {
            pad_print("TASK_RUNNING", 16);
        }
        else if (ptr->state == TASK_READY)
        {
            pad_print("TASK_READY", 16);
        }
        else if (ptr->state == TASK_BLOCKED)
        {
            pad_print("TASK_BLOCKED", 16);
        }

        sprintf(buf, "%p", ptr); // 调试用
        pad_print(buf, 16);

        putchk('\n');
    }
}